package binaryTree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import utils.binaryTreeUtil;
import utils.binaryTreeUtil.TreeNode;

public class lowestCommonAncestor {
    static HashMap<TreeNode, Integer> depMap = new HashMap<>();
    static HashMap<Integer, List<TreeNode>> map = new HashMap<>();

    // 先通过前序遍历，确定各个节点的深度
    public static void preOrder(TreeNode node, int depth) {
        if (node == null) {
            return;
        }
        depMap.put(node, depth);
        List<TreeNode> nodeList = map.getOrDefault(depth, new ArrayList<>());
        nodeList.add(node);
        map.put(depth, nodeList);
        preOrder(node.left, depth + 1);
        preOrder(node.right, depth + 1);
    }

    public static TreeNode getFather(int depth, TreeNode p, TreeNode q) {
        TreeNode pFather = null, qFather = null;
        while (depth >= 0) {
            List<TreeNode> list = map.get(depth - 1);
            for (TreeNode parentNode : list) {
                if (pFather != null && qFather != null) {
                    break;
                }
                if (parentNode.left == p || parentNode.right == p) {
                    pFather = parentNode;
                }
                if (parentNode.left == q || parentNode.right == q) {
                    qFather = parentNode;
                }
            }
            if (pFather!=null&&pFather == qFather) {
                return pFather;
            }
            p=pFather;
            q=qFather;
            pFather = null;
            qFather = null;
            depth--;
        }
        return pFather;
    }

    public static TreeNode lowestCommonAncestor(TreeNode root, TreeNode p, TreeNode q) {
        preOrder(root, 0);
        // 获取depth
        Integer pDep = depMap.get(p);
        Integer qDep = depMap.get(q);

        if (pDep == qDep) {
            // 寻找各自的父节点
            return getFather(pDep, p, q);
        } else if (pDep > qDep) {
            
            for (int i = pDep; i > qDep; i--) {
                // 找父节点
                TreeNode pFather = null;
                List<TreeNode> list = map.get(i - 1);
                for (TreeNode parentNode : list) {
                    if (pFather != null) {
                        break;
                    }
                    if (parentNode.left == p || parentNode.right == p) {
                        pFather = parentNode;
                        p=pFather;
                    }
                }
            }
            if (p==q) {
                return p;
            }
            return getFather(qDep, p, q);

        } else {
            for (int i = qDep; i > pDep; i--) {
                // 找父节点
                TreeNode qFather = null;
                List<TreeNode> list = map.get(i - 1);
                for (TreeNode parentNode : list) {
                    if (qFather != null) {
                        break;
                    }
                    if (parentNode.left == q || parentNode.right == q) {
                        qFather = parentNode;
                        q=qFather;
                    }
                }
            }
            if (p==q) {
                return p;
            }
            return getFather(pDep, p, q);
        }
    }

    public static void main(String[] args) {
        TreeNode root = binaryTreeUtil.buildTree(new Integer[]{3,5,1,6,2,0,8,null,null,7,4});
        lowestCommonAncestor(root,root.left.left,root.right.right);
    }
}
